#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
选择性数据导入脚本

用于从CSV文件中只导入指定的列（date、hour、type和站点ID列）

使用方法:
    python import_selected_columns.py --file [CSV文件路径] --sites [要保留的站点数量]
"""

import argparse
import os
import sys
import pandas as pd
import glob
from datetime import datetime
import logging

# 配置日志
logging.basicConfig(
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    level=logging.INFO
)
logger = logging.getLogger('import_selected_columns')

def import_csv_file(file_path, num_sites=10, output_file=None):
    """
    从CSV文件导入数据，只保留指定的列
    
    参数:
        file_path: CSV文件路径
        num_sites: 要保留的站点数量，默认为10
        output_file: 输出文件路径，如果为None则不保存
    
    返回:
        处理后的DataFrame
    """
    if not os.path.exists(file_path):
        logger.error(f"文件不存在: {file_path}")
        return None
    
    try:
        # 读取CSV文件
        logger.info(f"正在读取文件: {file_path}")
        df = pd.read_csv(file_path)
        logger.info(f"文件读取完成: {os.path.basename(file_path)}, {len(df)} 条记录")
        
        # 获取所有列名
        all_columns = df.columns.tolist()
        
        # 提取基础列和站点ID列
        base_columns = ['date', 'hour', 'type']
        # 站点ID列（所有不在基础列中的列）
        site_columns = [col for col in all_columns if col not in base_columns]
        
        # 只保留需要的列
        selected_site_columns = site_columns[:num_sites] if num_sites > 0 else []
        selected_columns = base_columns + selected_site_columns
        df_selected = df[selected_columns]
        
        # 打印结果
        logger.info(f"共选择了 {len(selected_columns)} 列:")
        logger.info(f"基础列: {base_columns}")
        logger.info(f"站点列: {selected_site_columns[:5]}{'...' if len(selected_site_columns) > 5 else ''}")
        logger.info(f"共 {len(df_selected)} 行数据")
        
        # 保存到文件（如果指定了输出文件）
        if output_file:
            df_selected.to_csv(output_file, index=False)
            logger.info(f"已保存到文件: {output_file}")
        
        return df_selected
    
    except Exception as e:
        logger.error(f"处理文件 {file_path} 失败: {str(e)}")
        return None

def import_csv_directory(csv_dir, num_sites=10, output_dir=None):
    """
    从目录中导入所有CSV文件，只保留指定的列
    
    参数:
        csv_dir: CSV文件目录
        num_sites: 要保留的站点数量，默认为10
        output_dir: 输出目录，如果为None则不保存
    
    返回:
        处理的文件数量
    """
    if not os.path.exists(csv_dir):
        logger.error(f"目录不存在: {csv_dir}")
        return 0
    
    # 获取所有CSV文件
    csv_files = glob.glob(os.path.join(csv_dir, "*.csv"))
    if not csv_files:
        logger.error(f"未找到CSV文件: {csv_dir}")
        return 0
    
    logger.info(f"找到 {len(csv_files)} 个CSV文件")
    
    # 创建输出目录（如果需要）
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir)
        logger.info(f"创建输出目录: {output_dir}")
    
    processed_files = 0
    for csv_file in csv_files:
        output_file = None
        if output_dir:
            output_file = os.path.join(output_dir, os.path.basename(csv_file))
        
        df = import_csv_file(csv_file, num_sites, output_file)
        if df is not None:
            processed_files += 1
    
    logger.info(f"总共处理了 {processed_files} 个文件")
    return processed_files

def main():
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='从CSV文件中只导入指定的列')
    parser.add_argument('--file', type=str, help='单个CSV文件路径')
    parser.add_argument('--dir', type=str, help='CSV文件目录')
    parser.add_argument('--sites', type=int, default=10, help='要保留的站点数量，默认为10')
    parser.add_argument('--output-file', type=str, help='输出文件路径（仅当指定--file时有效）')
    parser.add_argument('--output-dir', type=str, help='输出目录（仅当指定--dir时有效）')
    args = parser.parse_args()
    
    if args.file:
        # 处理单个文件
        import_csv_file(args.file, args.sites, args.output_file)
    elif args.dir:
        # 处理目录中的所有文件
        import_csv_directory(args.dir, args.sites, args.output_dir)
    else:
        logger.error("必须指定--file或--dir参数")
        parser.print_help()
        return 1
    
    return 0

if __name__ == "__main__":
    sys.exit(main())